{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "id": "ITN5GnDEeeZJ",
    "colab_type": "code",
    "colab": {}
   },
   "outputs": [],
   "source": [
    "# http://pytorch.org/\n",
    "from os.path import exists\n",
    "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
    "platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
    "cuda_output = !ldconfig -p|grep cudart.so|sed -e 's/.*\\.\\([0-9]*\\)\\.\\([0-9]*\\)$/cu\\1\\2/'\n",
    "accelerator = cuda_output[0] if exists('/dev/nvidia0') else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "xRIwwk0SeXNg",
    "colab_type": "code",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 233.0
    },
    "outputId": "2a4b4ba1-ddb7-4bc7-df24-1782844a11c6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1000 cost = 1.372447\n",
      "Epoch: 2000 cost = 0.822247\n",
      "Epoch: 3000 cost = 0.682714\n",
      "Epoch: 4000 cost = 1.896184\n",
      "Epoch: 5000 cost = 1.262253\n",
      "Epoch: 6000 cost = 0.943720\n",
      "Epoch: 7000 cost = 0.710601\n",
      "Epoch: 8000 cost = 0.431383\n",
      "Epoch: 9000 cost = 0.292406\n",
      "Epoch: 10000 cost = 0.196790\n",
      "Lorem ipsum dolor sit amet consectetur adipisicing elit sed do eiusmod tempor incididunt ut labore et dolore magna aliqua Ut enim ad minim veniam quis nostrud exercitation\n",
      "['ipsum', 'dolor', 'sit', 'amet', 'consectetur', 'adipisicing', 'elit', 'sed', 'do', 'eiusmod', 'eiusmod', 'incididunt', 'ut', 'labore', 'et', 'dolore', 'magna', 'aliqua', 'Ut', 'enim', 'ad', 'minim', 'veniam', 'quis', 'nostrud', 'exercitation']\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "  code by Tae Hwan Jung(Jeff Jung) @graykode\n",
    "'''\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "\n",
    "dtype = torch.FloatTensor\n",
    "\n",
    "sentence = (\n",
    "    'Lorem ipsum dolor sit amet consectetur adipisicing elit '\n",
    "    'sed do eiusmod tempor incididunt ut labore et dolore magna '\n",
    "    'aliqua Ut enim ad minim veniam quis nostrud exercitation'\n",
    ")\n",
    "\n",
    "word_dict = {w: i for i, w in enumerate(list(set(sentence.split())))}\n",
    "number_dict = {i: w for i, w in enumerate(list(set(sentence.split())))}\n",
    "n_class = len(word_dict)\n",
    "max_len = len(sentence.split())\n",
    "n_hidden = 5\n",
    "\n",
    "def make_batch(sentence):\n",
    "    input_batch = []\n",
    "    target_batch = []\n",
    "\n",
    "    words = sentence.split()\n",
    "    for i, word in enumerate(words[:-1]):\n",
    "        input = [word_dict[n] for n in words[:(i + 1)]]\n",
    "        input = input + [0] * (max_len - len(input))\n",
    "        target = word_dict[words[i + 1]]\n",
    "        input_batch.append(np.eye(n_class)[input])\n",
    "        target_batch.append(target)\n",
    "\n",
    "    return Variable(torch.Tensor(input_batch)), Variable(torch.LongTensor(target_batch))\n",
    "\n",
    "class BiLSTM(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BiLSTM, self).__init__()\n",
    "\n",
    "        self.lstm = nn.LSTM(input_size=n_class, hidden_size=n_hidden, bidirectional=True)\n",
    "        self.W = nn.Parameter(torch.randn([n_hidden * 2, n_class]).type(dtype))\n",
    "        self.b = nn.Parameter(torch.randn([n_class]).type(dtype))\n",
    "\n",
    "    def forward(self, X):\n",
    "        input = X.transpose(0, 1)  # input : [n_step, batch_size, n_class]\n",
    "\n",
    "        hidden_state = Variable(torch.zeros(1*2, len(X), n_hidden))   # [num_layers(=1) * num_directions(=1), batch_size, n_hidden]\n",
    "        cell_state = Variable(torch.zeros(1*2, len(X), n_hidden))     # [num_layers(=1) * num_directions(=1), batch_size, n_hidden]\n",
    "\n",
    "        outputs, (_, _) = self.lstm(input, (hidden_state, cell_state))\n",
    "        outputs = outputs[-1]  # [batch_size, n_hidden]\n",
    "        model = torch.mm(outputs, self.W) + self.b  # model : [batch_size, n_class]\n",
    "        return model\n",
    "\n",
    "input_batch, target_batch = make_batch(sentence)\n",
    "\n",
    "model = BiLSTM()\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# Training\n",
    "for epoch in range(10000):\n",
    "    optimizer.zero_grad()\n",
    "    output = model(input_batch)\n",
    "    loss = criterion(output, target_batch)\n",
    "    if (epoch + 1) % 1000 == 0:\n",
    "        print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "predict = model(input_batch).data.max(1, keepdim=True)[1]\n",
    "print(sentence)\n",
    "print([number_dict[n.item()] for n in predict.squeeze()])"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Bi-LSTM-Torch.ipynb",
   "version": "0.3.2",
   "provenance": [],
   "collapsed_sections": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "accelerator": "GPU"
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
